﻿using System;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;

/// <summary>
/// 适用于Unity Job System的原生优先队列实现
/// 基于最小堆数据结构，支持泛型元素和优先级
/// </summary>
/// <typeparam name="TElement">队列中存储的元素类型</typeparam>
/// <typeparam name="TPriority">用于排序的优先级类型</typeparam>
public struct NativePriorityQueue<TElement, TPriority> : IDisposable
    where TElement : unmanaged
    where TPriority : unmanaged, IComparable<TPriority>
{
    /// <summary>
    /// 优先队列中的节点，包含元素和对应的优先级
    /// </summary>
    private struct Node
    {
        /// <summary>
        /// 存储的元素
        /// </summary>
        public TElement Element;

        /// <summary>
        /// 用于排序的优先级值
        /// </summary>
        public TPriority Priority;

        /// <summary>
        /// 初始化节点的构造函数
        /// </summary>
        /// <param name="element">要存储的元素</param>
        /// <param name="priority">元素的优先级</param>
        public Node(TElement element, TPriority priority)
        {
            Element = element;
            Priority = priority;
        }
    }

    [NativeDisableContainerSafetyRestriction]
    private NativeArray<Node> nodes;

    private int count;
    private Allocator allocator;
    private bool isDisposed;

    /// <summary>
    /// 队列中的元素数量
    /// </summary>
    public int Count => count;

    /// <summary>
    /// 队列是否已创建（内存已分配）
    /// </summary>
    public bool IsCreated => nodes.IsCreated && !isDisposed;

    /// <summary>
    /// 初始化优先队列
    /// </summary>
    /// <param name="initialCapacity">初始容量</param>
    /// <param name="allocator">内存分配器</param>
    public NativePriorityQueue(int initialCapacity, Allocator allocator)
    {
        if (initialCapacity < 1)
            throw new ArgumentException("初始容量必须大于0", nameof(initialCapacity));

        this.allocator = allocator;
        nodes = new NativeArray<Node>(initialCapacity, allocator);
        count = 0;
        isDisposed = false;
    }

    /// <summary>
    /// 向队列中添加元素
    /// </summary>
    /// <param name="element">要添加的元素</param>
    /// <param name="priority">元素的优先级</param>
    public void Enqueue(TElement element, TPriority priority)
    {
        CheckDisposed();

        // 如果容量不足，扩容
        if (count == nodes.Length)
        {
            Resize();
        }

        // 添加元素并上浮调整堆
        nodes[count] = new Node(element, priority);
        SiftUp(count);
        count++;
    }

    /// <summary>
    /// 从队列中取出优先级最高的元素
    /// </summary>
    /// <returns>优先级最高的元素</returns>
    public TElement Dequeue()
    {
        CheckDisposed();

        if (count == 0)
        {
            throw new InvalidOperationException("无法从空队列中取出元素");
        }

        // 取出根节点（优先级最高）
        var top = nodes[0];
        count--;

        // 将最后一个元素移到根节点并下沉调整堆
        nodes[0] = nodes[count];
        SiftDown(0);

        return top.Element;
    }

    /// <summary>
    /// 查看队列中优先级最高的元素，但不移除它
    /// </summary>
    /// <returns>优先级最高的元素</returns>
    public TElement Peek()
    {
        CheckDisposed();

        if (count == 0)
        {
            throw new InvalidOperationException("队列为空");
        }

        return nodes[0].Element;
    }

    /// <summary>
    /// 清空队列中的所有元素
    /// </summary>
    public void Clear()
    {
        count = 0;
    }

    /// <summary>
    /// 释放队列占用的内存
    /// </summary>
    public void Dispose()
    {
        if (!isDisposed && nodes.IsCreated)
        {
            nodes.Dispose();
            isDisposed = true;
        }
    }

    /// <summary>
    /// 堆的上浮操作，确保堆的性质
    /// </summary>
    private void SiftUp(int index)
    {
        while (index > 0)
        {
            int parent = (index - 1) / 2;

            // 如果当前节点优先级高于父节点，停止上浮
            if (nodes[index].Priority.CompareTo(nodes[parent].Priority) >= 0)
                break;

            // 交换当前节点和父节点
            Swap(index, parent);
            index = parent;
        }
    }

    /// <summary>
    /// 堆的下沉操作，确保堆的性质
    /// </summary>
    private void SiftDown(int index)
    {
        int last = count - 1;

        while (true)
        {
            int left = 2 * index + 1;
            int right = 2 * index + 2;
            int smallest = index;

            // 找出当前节点、左子节点、右子节点中优先级最高的
            if (left <= last && nodes[left].Priority.CompareTo(nodes[smallest].Priority) < 0)
                smallest = left;

            if (right <= last && nodes[right].Priority.CompareTo(nodes[smallest].Priority) < 0)
                smallest = right;

            // 如果当前节点已是优先级最高的，停止下沉
            if (smallest == index)
                break;

            // 交换当前节点和优先级最高的子节点
            Swap(index, smallest);
            index = smallest;
        }
    }

    /// <summary>
    /// 交换两个节点的位置
    /// </summary>
    private void Swap(int i, int j)
    {
        (nodes[i], nodes[j]) = (nodes[j], nodes[i]);
    }

    /// <summary>
    /// 扩容操作，当容量不足时调用
    /// </summary>
    private void Resize()
    {
        int newCapacity = nodes.Length * 2;
        var newNodes = new NativeArray<Node>(newCapacity, allocator);

        // 复制现有元素到新数组
        for (int i = 0; i < count; i++)
        {
            newNodes[i] = nodes[i];
        }

        // 释放旧数组并替换为新数组
        nodes.Dispose();
        nodes = newNodes;
    }

    /// <summary>
    /// 检查队列是否已释放
    /// </summary>
    private void CheckDisposed()
    {
        if (isDisposed)
        {
            throw new ObjectDisposedException("NativePriorityQueue", "队列已被释放，无法执行操作");
        }
    }
}